package twilight.companion.entity.ai;

import net.minecraft.core.BlockPos;
import net.minecraft.tags.FluidTags;
import net.minecraft.world.entity.ai.goal.Goal;
import net.minecraft.world.entity.ai.navigation.PathNavigation;
import net.minecraft.world.level.pathfinder.Path;
import twilight.companion.entity.MiniYetiEntity;

import java.util.EnumSet;

public class MiniYetiSeekWaterGoal extends Goal {
    private final MiniYetiEntity yeti;
    private final double speedModifier;
    private BlockPos waterPos;
    private Path path;

    public MiniYetiSeekWaterGoal(MiniYetiEntity yeti, double speedModifier) {
        this.yeti = yeti;
        this.speedModifier = speedModifier;
        this.setFlags(EnumSet.of(Goal.Flag.MOVE));
    }

    @Override
    public boolean canUse() {
        if (!this.yeti.isOnFire()) {
            return false;
        }

        this.waterPos = this.findNearestWater();
        if (this.waterPos == null) {
            return false;
        }

        PathNavigation navigation = this.yeti.getNavigation();
        this.path = navigation.createPath(this.waterPos, 1);
        return this.path != null;
    }

    @Override
    public boolean canContinueToUse() {
        if (!this.yeti.isOnFire()) {
            return false;
        }
        
        if (this.waterPos == null) {
            return false;
        }

        double distanceToWater = this.yeti.distanceToSqr(
            this.waterPos.getX() + 0.5, 
            this.waterPos.getY(), 
            this.waterPos.getZ() + 0.5
        );

        if (distanceToWater <= 4.0) {
            return false;
        }
        return !this.yeti.getNavigation().isDone();
    }

    @Override
    public void start() {
        if (this.path != null) {
            this.yeti.getNavigation().moveTo(this.path, this.speedModifier);
        }
    }

    @Override
    public void stop() {
        this.waterPos = null;
        this.path = null;
        this.yeti.getNavigation().stop();
    }

    @Override
    public void tick() {
        if (this.yeti.getNavigation().isDone() && this.waterPos != null) {
            BlockPos newWaterPos = this.findNearestWater();
            if (newWaterPos != null && !newWaterPos.equals(this.waterPos)) {
                this.waterPos = newWaterPos;
                Path newPath = this.yeti.getNavigation().createPath(this.waterPos, 1);
                if (newPath != null) {
                    this.yeti.getNavigation().moveTo(newPath, this.speedModifier);
                }
            }
        }
    }

    private BlockPos findNearestWater() {
        BlockPos yetiPos = this.yeti.blockPosition();
        BlockPos bestWaterPos = null;
        double bestDistance = Double.MAX_VALUE;

        int searchRange = 16;
        for (int x = -searchRange; x <= searchRange; x++) {
            for (int y = -8; y <= 8; y++) {
                for (int z = -searchRange; z <= searchRange; z++) {
                    BlockPos checkPos = yetiPos.offset(x, y, z);
                    if (this.yeti.level().getFluidState(checkPos).is(FluidTags.WATER)) {
                        double distance = yetiPos.distSqr(checkPos);
                        if (distance < bestDistance) {
                            bestDistance = distance;
                            bestWaterPos = checkPos;
                        }
                    }
                }
            }
        }
        return bestWaterPos;
    }
}